SageMaker Processingで複数インスタンスを使う

#Amazon SageMaker

じょんすみす

2023.09.16

この記事は公開されてから1年以上経過しています。情報が古い可能性がありますので、ご注意ください。

おはこんハロチャオ～！何者(なにもん)なんじゃ？じょんすみすです。

SageMakerでは大規模なデータや大きなモデルの学習も行えるようにするために、複数インスタンスを使った分散処理ができるようになっています。今回は、Processingでインスタンス数を複数にした時の仕組みからその様子を見てみます。

複数インスタンスあるときに気に掛ける2つの要素

複数インスタンスを使う時に知っておくポイントとして大きく分けて2つあります。

動いている各インスタンスそのものに関する情報
各インスタンスで扱うデータについての制御

SageMaker Processingでは、これらについて扱うことが可能です。実際に必要な情報を元にどのような処理を行うかは実装する側に依存しますが、実装に必要な情報にアクセス可能になっているというわけです。

では、これらについて見ていきましょう。

動いている各インスタンスそのものに関する情報

まずは1つ目のインスタンスに関する情報です。

これらについて、SageMaker Processingでは以下のような情報が取得可能です。

動いているインスタンス数
各インスタンスのホスト名
自分がそれらのうちどれか

この情報はジョブを実行しているインスタンスの /opt/ml/config/resourceconfig.json に記述されています。内容としては以下のようになっています。

{
  "current_host": "algo-1",
  "hosts": ["algo-1","algo-2",...]
}

hostsに起動している全インスタンスのホスト名がリスト形式で記載されています。それぞれのホスト名を使ってインスタンス間の通信が可能ですし、このリストのサイズがインスタンス数となります。

それらのうち、自分自身がどれかはcurrent_hostに記載されています。こちらは、文字列型で単一のホスト名のみが記載されています。

自分の担当分の処理を決めたり、複数インスタンスでデータをやり取りしながら処理を行ったりすることが可能です。

各インスタンスで扱うデータについての制御

続いて、2つ目の要素について考えていきます。

SageMaker Processingでは、ProcessingInputs/ProcessingOutputsによってS3とローカルでデータをコピーして入出力を制御します。大規模なデータを扱うために複数インスタンスで処理を実行する場合に、全てのデータを単一のインスタンスにコピーされると都合が悪いですね。

ProcessingInputsでは ProcessingInputs/S3Input/S3DataDistributionType での指定でこれに対応します。この値は「FullyReplicated」「ShardedByS3Key」の2つから指定可能になっており以下のような動きをします。

FullyReplicated : 全てのインスタンスに全てのデータを渡す
ShardedByS3Key : インスタンス数に応じて均等割りしてそれぞれに別のデータを渡す

ShardedByS3Keyを指定すれば、均等かつ被りが無いようにデータを分割して各インスタンスにコピーしてくれます。ただし、データ全体の行数などまで見て分割するというわけではなく、ProcessingInputsで指定したプレフィックス以下のオブジェクト数で均等割りされます。そのため、複数オブジェクト存在していてそれぞれのサイズが大きく異なる場合などは各インスタンスでの処理負荷に偏りが生じる場合があるのでご注意ください。

S3DataDistributionTypeはインプットに指定している項目ごとに設定可能です。そのため、どのインスタンスでも使うマスターデータはFullyReplicated、各インスタンスで個別に処理する対象となるトランザクションデータはShardedByS3Keyといったように混在させた指定も可能です。利用方法に応じて指定を使い分けるのはユーザ側で考えることになりますが、柔軟な対応が可能です。